#coding=utf-8
import numpy as np
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
from matplotlib import pyplot as plt


tf.logging.set_verbosity(tf.logging.INFO)

mnist = input_data.read_data_sets("./", one_hot=True)

print(mnist.train.images.shape)
print(mnist.train.labels.shape)

print(mnist.validation.images.shape)
print(mnist.validation.labels.shape)

print(mnist.test.images.shape)
print(mnist.test.labels.shape)

plt.figure(figsize=(8, 8))

# for idx in range(16):
#     plt.subplot(4, 4, idx + 1)
#     plt.axis('off')
#     plt.title('[{}]'.format(np.argmax(mnist.train.labels[idx])))
#     plt.imshow(mnist.train.images[idx].reshape((28, 28)))

x = tf.placeholder("float", [None, 784], name='x')
y = tf.placeholder("float", [None, 10], name='y')

x_image = tf.reshape(x, [-1, 28, 28, 1])

with tf.name_scope('conv1'):
    C1 = tf.contrib.slim.conv2d(
        x_image, 6, [5, 5], padding='VALID', activation_fn=tf.nn.relu)

with tf.name_scope('pool1'):
    S2 = tf.contrib.slim.max_pool2d(C1, [2, 2], stride=[2, 2], padding='VALID')

with tf.name_scope('conv2'):
    C3 = tf.contrib.slim.conv2d(
        S2, 16, [5, 5], padding='VALID', activation_fn=tf.nn.relu)

with tf.name_scope('pool2'):
    S4 = tf.contrib.slim.max_pool2d(C3, [2, 2], stride=[2, 2], padding='VALID')

with tf.name_scope('fc1'):
    S4_flat = tf.contrib.slim.flatten(S4)
    C5 = tf.contrib.slim.fully_connected(
        S4_flat, 120, activation_fn=tf.nn.relu)

with tf.name_scope('fc2'):
    F5 = tf.contrib.slim.fully_connected(C5, 200, activation_fn=tf.nn.relu)

with tf.name_scope('fc3'):
    F6 = tf.contrib.slim.fully_connected(F5, 200, activation_fn=tf.nn.relu)

with tf.name_scope('dropout'):
    keep_prob = tf.placeholder(name='keep_prob', dtype=tf.float32)
    F6_drop = tf.nn.dropout(F6, keep_prob)

with tf.name_scope('fc4'):
    logits = tf.contrib.slim.fully_connected(F6_drop, 10, activation_fn=None)

cross_entropy_loss = tf.reduce_mean(
    tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=y))

l2_loss = tf.add_n([
    tf.nn.l2_loss(w)
    for w in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES)
])

for w in tf.get_collection(tf.GraphKeys.TRAINABLE_VARIABLES):
    print(w.name)
    tf.summary.histogram(w.name, w)

total_loss = cross_entropy_loss + 7e-5 * l2_loss
tf.summary.scalar('cross_entropy_loss', cross_entropy_loss)
tf.summary.scalar('l2_loss', l2_loss)
tf.summary.scalar('total_loss', total_loss)

learning_rate = tf.placeholder(tf.float32)

optimizer = tf.train.GradientDescentOptimizer(
    learning_rate=learning_rate).minimize(total_loss)

pred = tf.nn.softmax(logits)
correct_pred = tf.equal(tf.argmax(y, 1), tf.argmax(logits, 1))
accuracy = tf.reduce_mean(tf.cast(correct_pred, tf.float32))

batch_size = 100
trainig_step = 20000

saver = tf.train.Saver()

merged = tf.summary.merge_all()

lr = [0.3,0.2,0.1,0.05,0.02]

with tf.Session() as sess:

    writer = tf.summary.FileWriter("logs/", sess.graph)

    sess.run(tf.global_variables_initializer())

    #定义验证集与测试集
    validate_data = {
        x: mnist.validation.images,
        y: mnist.validation.labels,
        keep_prob: 1.0
    }
    test_data = {x: mnist.test.images, y: mnist.test.labels, keep_prob: 1.0}

    for i in range(trainig_step):
        xs, ys = mnist.train.next_batch(batch_size)
        _, loss, rs = sess.run(
            [optimizer, cross_entropy_loss, merged],
            feed_dict={
                x: xs,
                y: ys,
                keep_prob: 0.6,
                learning_rate: lr[int(i/(trainig_step/len(lr)))]
            })
        writer.add_summary(rs, i)

        #每100次训练打印一次损失值与验证准确率
        if i > 0 and i % 100 == 0:
            validate_accuracy = sess.run(accuracy, feed_dict=validate_data)
            print(
                "after %d training steps, the loss: %g, the validation: %g, learning rate: %s"
                % (i, loss, validate_accuracy,lr[int(i/(trainig_step/len(lr)))]))
    saver.save(sess, './model.ckpt')

    print("the training is finish!")
    #最终的测试准确率
    final_pred,acc = sess.run([pred,accuracy], feed_dict=test_data)
    print("the test accuarcy is:", acc)
    orders = np.argsort(final_pred)
    # plt.figure(figsize=(8, 8))
    for idx in range(16):
        order = orders[idx, :][-1]
        prob = final_pred[idx, :][order]
        plt.subplot(4, 4, idx + 1)
        plt.axis('off')
        plt.title('{}: [{}]-[{:.1f}%]'.format(
            np.argmax(mnist.test.labels[idx]), order, prob * 100))
        plt.imshow(mnist.test.images[idx].reshape((28, 28)))
    plt.show()